open Sast
open Type
open Operator

exception Bug of string (* For "impossible" situations *)

(* Main entry point: run a program *)
let translate prog out =
    (* Printing functions *)
    let put str = output_string out (str ^ " ") in
    let typestr = function
        | Float -> "float"
        | Int -> "integer"
        | Vector _ -> "vector"
        | RInt -> "randInt"
        | RFloat -> "randFloat"
        | List -> "list"
        | Str -> "string"
        | Tuple -> "tuple"
        | Func -> "function"
    in

    (* Conversion functions: wrap converters around expressions *)
    let rec extotp ex = function
        | Float -> extofloat ex
        | Int -> extoint ex
        | Str -> extostring ex
        | _ -> eval ex

    and extoint ex = 
        let _, tp = ex in
        match tp with
        | Int -> eval ex
        | Float -> put "(int_of_float"; eval ex; put ")"
        | RInt -> defrand ex
        | RFloat -> put "(int_of_float "; defrand ex; put ")"
        | _ -> raise (Bug ("Tried to convert "^(typestr tp)^" to integer"))

    and extofloat ex = 
        let _, tp = ex in
        match tp with
        | Int -> put "(float_of_int"; eval ex; put ")"
        | Float -> eval ex
        | RInt -> put "(float_of_int "; defrand ex; put ")"
        | RFloat -> defrand ex
        | _ -> raise (Bug ("Tried to convert "^(typestr tp)^" to float"))

    and extostring ex = 
        let _, tp = ex in
        match tp with
        | Str -> eval ex
        | Int -> put "(string_of_int"; eval ex; put ")"
        | Float -> put "(string_of_float"; eval ex; put ")"
        | RInt -> put "(string_of_int "; defrand ex; put ")"
        | RFloat -> put "(string_of_float "; defrand ex; put ")"
        | Vector _ -> (
            put "(let arr ="; eval ex;
            put ("in (\"<\" ^ string_of_float arr.(0) ^ " ^
                "(Array.fold_left (fun o f->(o^\",\"^(string_of_float f)))" ^
                "\"\" (Array.sub arr 1 (Array.length arr - 1)))^\">\"))")
            )
        | _ -> raise (Bug ("Tried to convert "^(typestr tp)^" to string"))

    and defrand ex =
        let _, tp = ex in
        match tp with
        | RInt -> put "(Rand.getRandInt "; eval ex; put ")"
        | RFloat -> put "(Rand.getRandFloat "; eval ex; put ")"
        | _ -> raise (Bug ("Tried to define a " ^ (typestr tp)))

    (* Evaluate an expression *)
    and eval = function

        (* Literals *)
        | LInt(i), _ -> put (string_of_int i)
        | LFloat(f), _ -> put (string_of_float f)
        | LList(l), tp -> (
            put "([";
            List.iter (fun ex -> (eval ex; put ";")) l;
            put "])")
        | LVctr(l), _ -> ( (* Vectors are arrays of floats *)
            put "([|";
            List.iter (fun ex -> extofloat ex; put ";") l;
            put "|])")
        | LTple(l), _ -> (* Tuples are tuples: have at least 2 elements *)
            (match l with
            | [] -> raise (Bug "Empty tuple")
            | hd::tl ->
                put "(";
                eval hd;
                List.iter ( fun ex -> put ","; eval ex) tl;
                put ")")
        | LString(s), _ -> put ("\"" ^ s ^ "\"")

        (* Identifiers.  Original names preserved with single quote *)
        | Id(var), _ -> put (var ^ "'")

        (* Array indices *)
        | Elmt(arr, pos), _ -> (
            put (arr ^ "'.(");
            let _, tp = pos in
            (match tp with
            | Int -> eval pos;
            | Float -> (
                put "int_of_float ";
                eval pos)
            | _ -> raise (Bug "Non-scalar array index"));
            put ")")

        (* Unary operators *)
        | Uop(op, ex), tp -> (
            put "(";
            (match op with
            | Minus -> (
                    match tp with
                    | Float -> (put "-."; extofloat ex)
                    | Int -> (put "-"; extoint ex)

                    (* Create a fresh array with inverted elements *)
                    | Vector n -> (
                        (* We need to evaluate the vector first, then create
                         * the inverted copy based off it *)
                        put "let arr = (";
                        eval ex;
                        put (") in Array.init " ^ string_of_int n
                            ^ " (fun i -> -. (arr.(i)))"))
                    | _ -> raise (Bug "Unary '-' expression is non-numeric")
                    )
            | Not -> (
                put "if ";
                let _, stp = ex in
                (match stp with
                | Int -> eval ex
                | Float -> ( put "(int_of_float"; eval ex; put ")")
                | _ -> raise (Bug "Unary '!' applied to non-scalar"));
                put "=0 then 1 else 0")
            );
            put ")")

        (* Binary operators.  This is big and (mostly) boring. *)
        | Binop(e1, op, e2), tp -> (
            let _, tp1 = e1 in
            put "(";
            (match tp with

            (* Integer result *)
            | Int -> (
                match op with

                (* Arithmetic *)
                | Add -> (extoint e1; put "+"; extoint e2)
                | Sub -> (extoint e1; put "-"; extoint e2)
                | Mult -> (extoint e1; put "*"; extoint e2)
                | Div -> (extoint e1; put "/"; extoint e2)
                | Rmndr -> (extoint e1; put "mod"; extoint e2)

                (* Comparison *)
                | Equal -> (put "if"; extofloat e1; put "="; extofloat e2; put "then 1 else 0")
                | Neq -> (put "if"; extofloat e1; put "<>"; extofloat e2; put "then 1 else 0")
                | Less -> (put "if"; extofloat e1; put "<"; extofloat e2; put "then 1 else 0")
                | Leq -> (put "if"; extofloat e1; put "<="; extofloat e2; put "then 1 else 0")
                | Greater -> (put "if"; extofloat e1; put ">"; extofloat e2; put "then 1 else 0")
                | Geq -> (put "if"; extofloat e1; put ">="; extofloat e2; put "then 1 else 0")

                (* Logical *)
                | Land -> (put "if ("; extoint e1; put "<> 0) && (";
                    extoint e2; put "<> 0) then 1 else 0")
                | Lor -> (put "if ("; extoint e1; put "= 0) && (";
                    extoint e2; put "= 0) then 0 else 1")

                | _ -> raise (Bug "Non-integer operator has integer type")
                )

            (* Float result *)
            | Float -> (
                match op with

                (* Arithmetic: 2 scalar operands *)
                | Add -> (extofloat e1; put "+."; extofloat e2)
                | Sub -> (extofloat e1; put "-."; extofloat e2)
                | Mult -> (extofloat e1; put "*."; extofloat e2)
                | Div -> (extofloat e1; put "/."; extofloat e2)

                (* Dot product *)
                | Innrp -> (
                    put "\nlet arr1 ="; eval e1; put "in";
                    put "\nlet arr2 ="; eval e2; put "in";
                    put "\nlet len1 = Array.length arr1 in";
                    put ("\nlet rec dot r = function -1-> r " ^ 
                    "| n-> dot (r+.arr1.(n)*.arr2.(n)) (n-1) in dot 0.0 (len1-1)")
                    )

                | _ -> raise (Bug "Non-float operator has float type")
                )

            (* Vector operators *)
            | Vector _ -> (
                (* We'll evaluate the vectors first, then apply the operator to
                 * the elements. *)
                (match tp1 with
                | Vector _ -> (put "\nlet arr1 = "; eval e1; put " in")
                | _ -> (put "\nlet flt1 = "; extofloat e1; put " in")
                );
                put "\nlet arr2 = "; eval e2; put " in";
                put "\nlet len1 = Array.length arr1 in";
                (* Check sizes? OCaml does it for me *)
                match op with

                (* Arithmetic: 2 vector operands *)
                | Add -> put "\nArray.init len1 (fun i-> arr1.(i)+.arr2.(i))"
                | Sub -> put "\nArray.init len1 (fun i-> arr1.(i)-.arr2.(i))"

                | Mult ->
                    (match tp1 with
                    (* Cross product.  Only valid for 3 dimension vectors *)
                    | Vector _ ->
                        put ("\n[| arr1.(1)*.arr2.(2) -. arr1.(2)*.arr2.(1) ; "
                        ^ "arr1.(2)*.arr2.(0) -. arr1.(0)*.arr2.(2) ; "
                        ^ "arr1.(0)*.arr2.(1) -. arr1.(1)*.arr2.(0) |]")

                    (* Vector Scaling, or 1x1 * 1xn matrix multiplication *)
                    | _ -> put "\nArray.init len1 (fun i-> flt1*.arr2.(i))"
                    )

                | _ -> raise (Bug "Non-vector operator has vector type")
                )

            (* List operators *)
            | List -> failwith "Lists not yet implemented"

            (* String operators *)
            | Str -> (
                match op with

                (* Concatenation *)
                | Add -> ( extostring e1; put "^"; extostring e2 )

                | _ -> raise (Bug "Non-string operator has string type")
                )

            | _ -> raise (Bug ("Binary operation on " ^ typestr tp))
            );
            put ")"
            )

        (* Function calls *)
        | Call(f, actuals), tp -> (
			let fname =
				if Builtin.exists f then
					Builtin.get_name f
				else
					(f ^ "'")
			in
            put ("(" ^ fname);
            (match actuals with 
            | [] -> ()
            | (ex,tp)::tl -> 
                    put "("; extotp ex tp; 
                    List.iter (fun (ex,tp) -> put ","; extotp ex tp) tl;
                    put ")");
            put ")")

        (* If expression *)
        | If(cond, truebody, falsebody), tp -> (
            put "(if";
            extoint cond;
            put "= 0 then (\n";
            extotp falsebody tp;
            put "\n) else (\n";
            extotp truebody tp;
            put "\n))")

        (* With expression *)
        | Scope(init, body), tp -> ( (* first run a declaration block, then body *)
            put "(";
            List.iter (fun stmt -> (dclr stmt; put "in\n")) init;
            extotp body tp;
            put ")")

    (* Print declarations *)
    and dclr = function
        | FDclr(tp, lname, formals, body) -> (
            put ("let rec " ^ lname ^ "'");
            (match formals with
            | [] -> ()
            | (_,name)::tl -> 
                    put ("("^name^"'"); 
                    List.iter (fun (_,name) -> put ","; 
                    put (name^"'")) tl;put ")");
            put "=";
            extotp body tp)
        | VDclr(tp, lname, actual) -> (
            put ("let " ^ lname ^ "' =");
            extotp actual tp)
    in

    (* Kickstart translation with a statement block, and keep an eye out for
     * "begin" (we need its type and its arguments'). *)
    let rec kick = function
    | [] -> raise (Bug "No begin function")
    | stmt::tl -> (
        dclr stmt;
        put ";;\n";
        (match stmt with
        | FDclr(beg_tp, name, arg_ls, _) ->
                if name = "begin" then
                    (* Anything after begin can't be reached, anyway *)
                    beg_tp, List.fold_left (fun l (t,_)->t::l) [] arg_ls
                else
                    kick tl
        | _ -> kick tl)
    )
    in kick prog

